#include "kernel_operator.h"
#include <cmath>
#include "math.h"

using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2;      //昇腾双buffer技术

class KernelLogSumExp {
public:
    __aicore__ inline KernelLogSumExp() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y,int32_t size, int32_t x_ndarray[], int32_t x_dimensional, int64_t x_dim[], int32_t dim_dimensional, bool keep_dim)
    {

        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");


        this->x_dimensional = x_dimensional;  //x的维度
        for(int i = 0; i < x_dimensional; i++)
        {
            this->x_ndarray[i] = x_ndarray[i];  //每个维度的元素个数
        }
        this->size =size;                     //总元素个数

        this->dim_dimensional = dim_dimensional;  //axis的个数
        for(int i = 0; i<this->dim_dimensional; i++)  
        {
            this->dim[i] = x_dim[i];
            if(this->dim[i] < 0)
                this->dim[i] += this->x_dimensional;
        }
		//printf("size is %d\n", this->size);
		//printf("dim_dimensional is %d", this->dim_dimensional);
        xGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_X *>(x), size);  //+8
        yGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_Y *>(y), size);  //+8

        int32_t cycles = 1;
        int32_t interval = 1;
        int32_t loopCount = 1;

        for(int i = 0; i < this->x_dimensional; i++)   //如果是0，循环的就是第1维，1，1+2维，2，就是1+2+3维的总元素个数
        {
            loopCount *= this->x_ndarray[i];
        }

        for(int i = 0; i<this->dim_dimensional; i++)
        {
            loopCount = loopCount / this->x_ndarray[this->dim[i]];
        }

        for(int i = 0; i<this->dim_dimensional; i++)
        {
            cycles *= this->x_ndarray[this->dim[i]];
        }

        
        for(int32_t i = this->dim[this->dim_dimensional-1]+1; i < this->x_dimensional; i++)
        {
            interval *= this->x_ndarray[i];     //跨其他维度时的跨度
        }

        printf("cycle is %d, interval is %d, loopCount is %d\n",cycles, interval, loopCount);
        this->cycles = cycles;
        this->interval = interval;
        this->loopCount = loopCount;
        printf("cycle is %d, interval is %d, loopCount is %d\n",this->cycles, this->interval, this->loopCount);

        // 这里的分块逻辑是8192个x,1024个y和buff,也许不是最好，先用着
        // 注意这里是interval为1，也就是连续才能并行，否则串行
        if(this->interval == 1)   
        {
            pipe.InitBuffer(inQueueX, BUFFER_NUM, 8192 * sizeof(DTYPE_X));
            pipe.InitBuffer(outQueueY, BUFFER_NUM, 8192 * sizeof(DTYPE_Y));
            pipe.InitBuffer(QueueBuff, 8192 * sizeof(DTYPE_X));
        }

    }

    // 计算阶乘
    __aicore__ float factorial(int n) {
        float result = 1.0f;
        for (int i = 1; i <= n; i++) {
            result *= i;
        }
        return result;
    }

    // 计算 e^r 的泰勒展开
    __aicore__ float exp_taylor(float r, int terms) {
        float result = 1.0f;
        float power = 1.0f;
        for (int i = 1; i < terms; i++) {
            power *= r;
            result += power / factorial(i);
        }
        return result;
    }

    // 精确计算 e^x
    __aicore__ float exp_precise(float x) {
        // 特殊值处理
        if (x > 88.0f) return HUGE_VALF;  // 溢出（float精度下的最大值约为 88）
        if (x < -103.0f) return 0.0f;    // 下溢（float精度下的最小值约为 -103）
        if (x == 0) return 1.0f;

        // 分解 x = n + r
        int n = (int)x;                 // 整数部分
        float r = x - n;                // 小数部分

        // 计算 e^n
        float exp_n = 1.0f;
        if (n > 0) {
            exp_n = 1.0f;
            for (int i = 0; i < n; i++) {
                exp_n *= 2.718281828459045f;  // 近似e
            }
        } else if (n < 0) {
            exp_n = 1.0f;
            for (int i = 0; i > n; i--) {
                exp_n /= 2.718281828459045f;  // 近似e
            }
        }

        // 计算 e^r
        float exp_r = exp_taylor(r, 10); // 使用 10 项泰勒展开

        // 返回结果
        return exp_n * exp_r;
    }


    // 多项式逼近计算 ln(1 + y)
    __aicore__ float log_approx(float y) {
        float result = y;
        float term = y;
        for (int n = 2; n <= 10; ++n) { // 使用10项展开
            term *= -y; // y^n
            result += term / n; // 加上 y^n/n
        }
        return result;
    }

    // 精确计算 ln(x)
    __aicore__ float log_precise(float x) {
        if (x <= 0.0f) {
            printf("Error: Input must be positive.\n");
            return -INFINITY; // 定义域错误
        }
        if (x == 1.0f) {
            return 0.0f; // 特殊值处理
        }

        int exponent = 0;

        // 将 x 分解为 m * 2^n，确保 m ∈ [1, 2)
        while (x > 1.41421356237f) {
            x /= 2.0f;
            exponent++;
        }
        while (x < 0.70710678118f) {
            x *= 2.0f;
            exponent--;
        }

        // 计算 ln(m) 使用多项式逼近
        float m = x - 1; // 转化为 ln(1 + m)
        float ln_m = log_approx(m);

        // 直接使用 ln(2) 进行计算
        return exponent * 0.69314718056f + ln_m; // ln(2) ≈ 0.69314718
    }

    __aicore__ inline void Process()
    {
        // if constexpr (std::is_same_v<DTYPE_Y, float>)
        // {
        //     printf("is_same float\n");
        //     if((this->interval == 1) && (8192%this->cycles == 0) && (this->cycles <= 8192))
        //     {
        //         printf("this->interval == 1\n");
        //         int capacity = 8192/this->cycles;    //这个设计就有点离谱，必须是8192能除尽的数才能并行,this->cycles不能为空
        //         printf("capacity is %d\n",capacity);
        //         for(int32_t i = 0; i < this->loopCount;)
        //         {
        //             if((this->loopCount - i) >= capacity)
        //             {
        //                 CopyIn(i, capacity);
        //                 Compute(i, capacity);
        //                 CopyOut(i, capacity);
        //             }
        //             else
        //             {
        //                 CopyIn(i, this->loopCount - i);
        //                 Compute(i, this->loopCount - i);
        //                 CopyOut(i, this->loopCount - i);
        //             }
        //             i+=capacity;
        //         }
        //     }
        //     else
        //     {
        //         //公式：b = x.max()
        //         //return b + np.log(np.sum(np.exp(x - b)))

        //         printf("this->interval != 1\n");
        //         float temp_sum;
        //         float temp_add;
        //         float temp_max;
        //         for(int z = 0; z < this->loopCount; z++)
        //         {
        //             int32_t x_num = z/this->interval;
        //             x_num = x_num*this->cycles*this->interval + z%this->interval;
        //             //max
        //             for(int i = 0; i < this->cycles; i++)
        //             {
        //                 int32_t temp_index = x_num + i*this->interval;
        //                 temp_add = xGm.GetValue(temp_index);
                        
        //                 if(i == 0)
        //                 {
        //                     temp_max = temp_add;
        //                 }
        //                 else
        //                 {
        //                     if(temp_max < temp_add )
        //                         temp_max = temp_add;
        //                 }
        //             }
        //             printf("temp_max is %f\n",temp_max);
        //             //exp
        //             for(int i = 0; i < this->cycles; i++)
        //             {
        //                 int32_t temp_index = x_num + i*this->interval;
        //                 temp_add = xGm.GetValue(temp_index);
        //                 printf("bef temp_add is %f\n",temp_add);
        //                 temp_add = temp_add - temp_max;
        //                 temp_add = exp_precise(temp_add);
        //                 xGm.SetValue(temp_index,temp_add);
        //                 printf("exp temp_add is %f\n",temp_add);
                        
        //                 //xGm.SetValue(temp_index,(half)exp((float)temp_add-(float)temp_max));
        //             }
        //             //sum
        //             for(int i = 0; i < this->cycles; i++)
        //             {
        //                 int32_t temp_index = x_num + i*this->interval;
        //                 temp_add = xGm.GetValue(temp_index);
                        
        //                 if(i == 0)
        //                 {
        //                     temp_sum = temp_add;
        //                 }
        //                 else
        //                 {
        //                     temp_sum = temp_sum + temp_add;
        //                 }
        //             }
        //             printf("temp_sum is %f\n",temp_sum);

        //             //max + log(sum)
        //             yGm.SetValue(z,(DTYPE_Y)(temp_max + log_precise(temp_sum)));
        //         }
        //     }
        // }
        // else
        {
            printf("is same half\n");
            float temp_sum;
            float temp_add;
            float temp_max;
            for(int z = 0; z < this->loopCount; z++)
            {
                int32_t x_num = z/this->interval;
                x_num = x_num*this->cycles*this->interval + z%this->interval;
                //max
                for(int i = 0; i < this->cycles; i++)
                {
                    int32_t temp_index = x_num + i*this->interval;
                    temp_add = (float)xGm.GetValue(temp_index);
                    
                    if(i == 0)
                    {
                        temp_max = temp_add;
                    }
                    else
                    {
                        if(temp_max < temp_add )
                            temp_max = temp_add;
                    }
                }
                printf("temp_max is %f\n",temp_max);
                //exp
                for(int i = 0; i < this->cycles; i++)
                {
                    int32_t temp_index = x_num + i*this->interval;
                    temp_add = (float)xGm.GetValue(temp_index);
                    printf("bef temp_add is %f\n",temp_add);
                    temp_add = temp_add - temp_max;
                    temp_add = exp_precise(temp_add);
                    printf("exp temp_add is %f\n",temp_add);
                    xGm.SetValue(temp_index,(DTYPE_Y)temp_add);
                    
                    //xGm.SetValue(temp_index,(half)exp((float)temp_add-(float)temp_max));
                }
                //sum
                for(int i = 0; i < this->cycles; i++)
                {
                    int32_t temp_index = x_num + i*this->interval;
                    temp_add = (float)xGm.GetValue(temp_index);
                    
                    if(i == 0)
                    {
                        temp_sum = temp_add;
                    }
                    else
                    {
                        temp_sum = temp_sum + temp_add;
                    }
                }
                printf("temp_sum is %f\n",temp_sum);
                //max + log(sum)
                yGm.SetValue(z,(DTYPE_Y)(temp_max + log_precise(temp_sum)));
            }
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t i, int32_t capacity)
    {
        LocalTensor<DTYPE_X> xLocal = inQueueX.AllocTensor<DTYPE_X>();
        DataCopy(xLocal, xGm[i*this->cycles], (this->cycles*capacity+7)/8*8);
        inQueueX.EnQue(xLocal);

    }
    __aicore__ inline void Compute(int32_t i, int32_t capacity)
    {
        LocalTensor<float> yLocal = outQueueY.AllocTensor<float>();

        LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        auto buff = QueueBuff.Get<float>();

        //max
        float tmp = 0.0;
        for(int z=0;z<capacity;z++)
        {
            ReduceMax(yLocal[z], xLocal[this->cycles*z], buff, this->cycles); 
            tmp = (float)yLocal[z].GetValue(0);
            //printf("max is %f\n",tmp);
            Duplicate(yLocal[this->cycles*z], tmp, this->cycles);
        }
        Sub(xLocal,xLocal,yLocal,8192);
        //printf("sub is %f\n",xLocal(0));
        
        //exp
        Exp(xLocal,xLocal,8192);
        //printf("exp is %f\n",xLocal(0));
        for(int z=0;z<capacity;z++)
        {
            ReduceSum(xLocal[z], xLocal[this->cycles*z], buff, this->cycles);  //sum
            ReduceMax(yLocal[z], yLocal[this->cycles*z], buff, this->cycles);  //max
        }
                    
        //max + log(sum)
        Ln(xLocal,xLocal,capacity);
        Add(yLocal,xLocal,yLocal,capacity);
        //printf("max + log(sum) is %f\n",yLocal(0));
        
        outQueueY.EnQue<float>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t i, int32_t j)
    {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        //DataCopy(yGm[i], yLocal, (j+7)/8*8);
        DataCopy(yGm[i], yLocal, (j+7)/8*8);
        //printf("yGm %d is %f\n",i, yGm(i));
        outQueueY.FreeTensor(yLocal);
    }


private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueBuff;
    
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;

    int32_t size;
    int32_t x_ndarray[20];
    int32_t x_dimensional;
    int64_t dim[20];
    int32_t dim_dimensional;
    bool keep_dim;

    int32_t cycles;
    int32_t interval;
    int32_t loopCount;
};

extern "C" __global__ __aicore__ void log_sum_exp(GM_ADDR x, GM_ADDR y,  GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelLogSumExp op;
    //补充init和process函数调用内容
    op.Init(x, y,tiling_data.size, tiling_data.x_ndarray, tiling_data.x_dimensional, tiling_data.x_dim,tiling_data.dim_dimensional,tiling_data.keep_dim);
    op.Process();
}

